
import os
import pickle
import random

import numpy as np
import copy

import torch
import time

from .utils import state_to_features, ACTIONS, vaild_action_filter
from .model import Framework

def setup(self):
    """
    Setup your code. This is called once when loading each agent.
    Make sure that you prepare everything such that act(...) can be called.

    When in training mode, the separate `setup_training` in train.py is called
    after this method. This separation allows you to share your trained agent
    with other students, without revealing your training code.

    In this example, our model is a set of probabilities over actions
    that are is independent of the game state.

    :param self: This object is passed to all callbacks and you can set arbitrary values.
    """
    # if self.train or not os.path.isfile("my-saved-model.pt"):
    if self.train:
        self.logger.info("Setting up model from scratch.")
        self.framework = Framework()
    else:
        self.logger.info("Loading model from saved state.")
        self.framework = Framework()
        self.framework.load_model('my-saved-model.pt')

def act(self, game_state: dict) -> str:
    """
    Your agent should parse the input, think, and take a decision.
    When not in training mode, the maximum execution time for this method is 0.5s.

    :param self: The same object that is passed to all of your callbacks.
    :param game_state: The dictionary that describes everything on the board.
    :return: The action to take as a string.
    """
    action_ids = vaild_action_filter(game_state)
    if random.random() > 0.2:
        observation = state_to_features(game_state) # [12, 17, 17]的观察矩阵
        action_score = self.framework.predict(observation)
        max_id = 0
        max_score = -10000000
        for aid in action_ids:
            if action_score[aid] > max_score:
                max_score = action_score[aid]
                max_id = aid
    else:
        max_id = random.choice(action_ids)

    self.logger.debug(f"Action {ACTIONS[max_id]}.")
    return ACTIONS[max_id]

